package 力扣.排序.快速排序;

public class 最小的k个数Offer40 {
    /**
     * 1、寻找在数组第k个索引位置的元素（实际是第k+1小的元素）
     *   如何做？使用快排查找TOPK
     * 2、遍历数组，数组元素和刚才找的元素进行比较。
     *            情况1：小于：则将其放入结果数组
     *            情况2：等于：记录相等元素的数量
     * 3、若结果数组中元素个数不够，则肯定在相等的元素里面，将相等的元素放入结果数组，直到填满数组。
     * @param arr
     * @param k
     * @return
     */
    public int[] getLeastNumbers(int[] arr, int k) {
        final int preKey = kth(arr, 0, arr.length, k);//寻找在数组第k个索引位置的元素（实际是第k+1小的元素）
        int[] ans = new int[k];
        int kid = 0;//记录ans数组中元素的数量
        int equal = 0;//记录相等元素的数量
        for (int te: arr) {
            if (te < preKey){
                ans[kid++] = te;
            }
            if (te == preKey){
                equal++;
            }
        }
        while (kid < k && equal-- > 0){
            ans[kid++] = preKey;
        }
        return ans;
    }
    private int kth(int[] nums,int b,int e,int k){
        if (b >= e){
            return 0;
        }
        if (b + 1 >= e){
            return nums[b];
        }

        int mid = nums[b + ((e - b) >> 1)];
        int l = b,i = b,r = e -1;
        while (i <= r){
            if (nums[i] == mid){
                i++;
            }else if(nums[i] < mid){
                swap(nums,i++,l++);
            }else{
                swap(nums,i,r--);
            }
        }
        int lCnt = l - b;
        int midCnt = i - l;
        if (k < lCnt){
            return kth(nums, b, l, k);
        }else if(k > lCnt + midCnt){
            return kth(nums, i, e, k - lCnt - midCnt);
        }else{
            return mid;
        }
    }

    private static void swap(int[] A, int i, int j) {
        int t = A[i];
        A[i] = A[j];
        A[j] = t;
    }
}
